#include "codegen.hpp"
// #define GEN_DEBUG
namespace backend{
    using RegId = RegManager::RegId;
    // int frame_size = 0;

    string CodeGen::code_gen(){
        string code;
        code += module_gen(compUnit);
        return code;
    }

    string CodeGen::module_gen(CompUnit *module){
        string code;
        string dataCode;
        string textCode;
        //clear last module's label record
        clearModuleRecord(module);
        //generate asmcode for all global values 
        SymbolTable::ModuleScope scope(compUnit->symbols);
        dataCode += globaldata_gen();

        code += space + ".arch armv7ve " + endl;
        code += space + ".text " + endl;
        auto functions = module->getFunctions();
        
        for(auto iter = functions->begin(); iter != functions->end(); ++iter){
            string name = iter->first;
            Function *func = iter->second;
            auto bblist = func->getBasicBlocks();
            if(bblist.empty()) continue;
            //generate asmcode for each function
            textCode += function_gen(func) + endl;
        }
        code += (dataCode + textCode + endl);
        return code;
    }

    string CodeGen::functionHead_gen(Function *func){
        string code;
        code += space + ".globl " + func->getName() + endl;
        code += space + ".p2align " + std::to_string(int_p2align) + endl;
        code += space + ".type " + func->getName() + ", %function" + endl;
        code += func->getName() + ":" + endl;
        return code;
    }
    /**
     * stack structure:
     * 
     * last function stack
     * --------------------------   <-- fp  point
     * callee preserved regs (include fp, sp, lr)
     * --------------------------
     * tempory variable(ir inst)
     * --------------------------
     * arg0,arg1,arg2,arg3(as tempory variable)
     * --------------------------    
     * dynamic stack(by AllocaInst)  ********
     * -------------------------- 
     * caller preserved regs
     * --------------------------
     * arg4 ~ argX
     * --------------------------    <-- sp point
     * next function stack
     * 
    */
    /**
     * prologue :
     *          preserve callee-saved register (lr, fp and other callee-saved regs)
     *          set new fp
     *          alloc stack space for local var/ callInstr / return value.
     *          store callInstr to stack
     * */
    string CodeGen::prologueCode_gen(Function *func, int *frame_size){
        string code;
        /** 
         *code in here
        */
        // return code;
        // LR: r14 + FP: r11 + other callee-save regs[r4 ~ r10]
        //
        // frame_size += 4 + 4;
        // r4 ~ r10
        // frame_size += (10 - 4 + 1) * 4;   
        code += M_emitComment("now frame_size = " + std::to_string(*frame_size));
        // code += space + "push" + "\t" + "{lr}" + endl;
        code += space + "push" + " " + "{fp, lr}" + endl;
        // code += space + "str" + "\t" + "fp, [sp, #-4]";
        code += space + "add" + " " + "fp, sp, #4" + endl;

        // 目前不考虑保存r4-r10，原因是所有变量考虑保存至栈空间中
        auto entryblock = func->getEntryBlock();
        // for (auto &arg : entryblock->getArguments()) {
        //     frame_size += getSize(arg->getType());
        // }
        auto bbs = func->getBasicBlocks();
        // for (auto iter = bbs.begin(); iter != bbs.end(); ++iter) {
        //     auto bb = iter->get();
        //     for(auto &instr : bb->getInstructions()){
        //         auto instrType = instr->getKind();
        //         // code += instruction_gen(instr.get());

        //         // pass func's subfunc callInstr
        //         if (instrType == Instruction::kCall) {
        //             auto arg_num = 0;
        //             auto callInstr = dynamic_cast<CallInst*>(instr.get());
        //             for (auto arg: callInstr->getArguments()) {
        //                 arg_num += 1;
        //                 if (arg_num > 4)
        //                     frame_size += getSize(arg);
        //             }
        //         } 
        //         // func's local variables
        //         // if (instrType == Instruction::kAlloca) {
        //         //     auto ops = instr->getType();
        //         //     frame_size += getSize(ops);
        //         // }
        //     }
        // } 

        // int nowoff = 8;

        // 该函数的参数 符号表记录
        int now_arg_num = 0;
        for (auto &arg : entryblock->getArguments()) {
            *frame_size += getSize(arg->getType());
            compUnit->symbols.insert(arg->getName(), (Value*)(arg.get()), *frame_size);
            code += space + "// add symbols as param: " + arg->getName() + " "
                            + std::to_string(*frame_size) + endl;
            if (now_arg_num < 4) {
                code += space + "str r" + std::to_string(now_arg_num) + ", "
                              + "[fp, #-" + std::to_string(*frame_size) + "]" + endl;
            } else {
                // TODO add subfunc params' pass
            }
            ++now_arg_num;
        }
        // 该函数的申请空间 符号表记录
        
        // for (auto iter = bbs.begin(); iter != bbs.end(); ++iter) {
        //     auto bb = iter->get();
        //     for (auto &instr : bb->getInstructions()) {
        //         auto instrType = instr->getKind();
        //         if (instrType == Instruction::kAlloca) {
        //             auto ops = instr->getType();
        //             nowoff += getSize(ops);
        //             code += space + "// add symbols as alloca: " + instr->getName() + " "
        //                             + std::to_string(frame_size - nowoff) + endl;   
        //             compUnit->symbols.updateOffset(instr->getName(), frame_size - nowoff);
        //         }

        //     }
        // }
        // 该函数调用子函数的传递参数 符号表记录
        for (auto iter = bbs.begin(); iter != bbs.end(); ++iter) {
            auto bb = iter->get();
            for (auto &instr : bb->getInstructions()) {
                auto instrType = instr->getKind();
                if (instrType == Instruction::kCall) {
                    auto callInstr = dynamic_cast<CallInst*>(instr.get());
                    auto arg_num = callInstr->getNumOperands();
                    for (auto iter = arg_num - 1; iter > 0; --iter) {
                        auto arg = callInstr->getOperand(iter);
                        arg_num -= 1;
                        if (arg_num > 4) {
                            *frame_size += getSize(arg);
                            compUnit->symbols.insert(arg->getName(), (Value*)(arg), *frame_size);
                            code += space + "// add symbols as pass arg: " + arg->getName() + " "
                                          + std::to_string(*frame_size) + endl;   
                                        //   + std::to_string(compUnit->symbols.getOffset(arg->getName())) + endl; 
                        }
                    }
                }
            }
        }
        
        
        code += space + "sub sp, sp, #" + std::to_string(*frame_size) + endl;

        return code;
   }

    /* epilogue : 
     *          free stack space
     *          restore sp
     *          restore used callee-saved register(lr, fp and other callee-saved regs)
     *          bx lr
    */
    string CodeGen::epilogueCode_gen(Function *func){
        string code;
        /** 
         *code in here
        */
        return code;
    }

    string CodeGen::function_gen(Function *func){
        curFunc = func;
        clearFunctionRecord(func);
        string bbCode;
        SymbolTable::FunctionScope scope(compUnit->symbols);
        auto bbs = func->getBasicBlocks();
        int frame_size = 4 + 4;
        for(auto iter = bbs.begin(); iter != bbs.end(); ++iter){
            auto bb = iter->get();
            bbCode += basicBlock_gen(bb, &frame_size);
        }
        string code;
        string funcHead = functionHead_gen(func);
        string prologueCode = prologueCode_gen(func, &frame_size);
        string epilogueCode = epilogueCode_gen(func);
        string literalPoolsCode = literalPoolsCode_gen(func);
        //
        code = funcHead + prologueCode  + bbCode + \
                    epilogueCode + literalPoolsCode;
        return code;
    }

    string CodeGen::basicBlock_gen(BasicBlock *bb, int *frame_size){
        curBB = bb;
        string bbLabel = getBBLabel(bb);
        SymbolTable::BlockScope scope(compUnit->symbols);
        string code;
        code += bbLabel + ":" + endl;
        for(auto &instr : bb->getInstructions()){
            auto instrType = instr->getKind();
            code += instruction_gen(instr.get(), frame_size);
        }
        return code;
    }
    /**
     * RegId : binaryInst_gen returns RegId as its destination operand
     * code  : asmcode generated by binaryInst_gen
    */
    pair<RegId, string> CodeGen::binaryInst_gen(BinaryInst *bInst, RegId dstRegId){
        string code;
        /** 
         *code in here
        */
        Value* lhs = bInst->getLhs();
        Value* rhs = bInst->getRhs();
        string bInstName;
        switch (bInst->getKind())
        {
        case Instruction::kAdd:
            bInstName = "add";
            break;
        case Instruction::kSub:
            bInstName = "sub";
            break;
        case Instruction::kMul:
            bInstName = "mul";
            break;
        case Instruction::kDiv:
            bInstName = "div";
            break;
        case Instruction::kRem:
            bInstName = "rem";
            break;
        case Instruction::kLOr:
            bInstName = "orr";
            break;
        case Instruction::kLAnd:
            bInstName = "and";
            break;
        case Instruction::kICmpEQ:
        case Instruction::kICmpNE:
        case Instruction::kICmpLT:
        case Instruction::kICmpGT:
        case Instruction::kICmpLE:
        case Instruction::kICmpGE:
        case Instruction::kFAdd:
            bInstName = "fadd";
            break;
        case Instruction::kFSub:
            bInstName = "fsub";
            break;
        case Instruction::kFMul:
            bInstName = "fmul";
            break;
        case Instruction::kFDiv:
            bInstName = "fdiv";
            break;
        case Instruction::kFRem:
            bInstName = "frem";
            break;
        case Instruction::kFCmpEQ:
        case Instruction::kFCmpNE:
        case Instruction::kFCmpLT:
        case Instruction::kFCmpGT:
        case Instruction::kFCmpLE:
        case Instruction::kFCmpGE:
        default:
            break;
        }
        code += "\t" + bInstName;
        code += " r" + bInst->getName() + ",";
        code += " r" + lhs->getName() + ",";
        code += " r" + rhs->getName() + "\n";
        code += M_emitComment("binary inst");
        
    
        return {dstRegId, code};
    }

    pair<RegId, string> CodeGen::unaryInst_gen(UnaryInst *uInst, RegId dstRegId){
        string code;
        /** 
         *code in here
        */
        Value* operand = uInst->getOperand();
        string uInstName;
        switch(uInst->getKind())
        {
        case  Instruction::kNeg:
            uInstName = "neg";
            break;
        case  Instruction::kNot:
            uInstName = "not";
            break;
        case  Instruction::kFNeg:
            uInstName = "fneg";
            break;
        case  Instruction::kFtoI:
            uInstName = "ftoi";
            break;
        case  Instruction::kIToF:
            uInstName = "itof";
            break;
        default:
            break;    
        }
        code += "\t" + uInstName;
        code += " r" + uInst->getName() + ",";
        code += " r" + operand->getName() + "\n";
        code += M_emitComment("unary inst");

        return {dstRegId, code};
    }
    pair<RegId, string> 
    CodeGen::allocaInst_gen(AllocaInst *aInst, RegManager::RegId dstRegId, int *frame_size){
        string code;
        /** 
         *code in here
        */

        auto ops = aInst->getType();
        *frame_size += getSize(ops);
        code += space + "// add symbols as alloca: " + aInst->getName() + " "
                        + std::to_string(*frame_size) + endl;   
        // std::cout << code;
        // std::cout.flush();
        compUnit->symbols.insert(aInst->getName(), (Value *)(aInst), *frame_size);
        return {dstRegId, code};
    }

    string CodeGen::storeInst_gen(StoreInst *stInst){
        string code;
        /** 
         *code in here
        */
        auto reg = stInst->getOperand(0);
        auto arg = stInst->getOperand(1);
        // code += space + "// store " + reg->getName() + ", " + arg->getName() + endl;
        code += space + "str, r" + stInst->getName() + ", [fp, #-" 
              + std::to_string(compUnit->symbols.getOffset(arg->getName())) 
              + "]" + endl;
        return code;
    }
    pair<RegId, string> 
    CodeGen::loadInst_gen(LoadInst  *ldInst, RegId dstRegId){
        string code;
        /** 
         *code in here
        */
        auto reg = ldInst->getOperand(0);
        // for (auto it = reg.begin(); it != reg.end(); ++it) {
        //     std::cout << it->getName() << "\n";
        //     std::cout.flush();
        // }
        // auto arg = ldInst->getOperand(1);
        code += space + "load r" + ldInst->getName() + ", [fp, #-"
              + std::to_string(compUnit->symbols.getOffset(reg->getName())) + "]" + endl;
        
        return {dstRegId, code};
    }
    string CodeGen::returnInst_gen(ReturnInst *retInst){
        string code;
        /** 
         *code in here
        */
        return code;
    }
    string CodeGen::uncondBrInst_gen(UncondBrInst *ubInst){
        string code;
        /** 
         *code in here
        */
        return code;
    }
    string CodeGen::condBrInst_gen(CondBrInst *cbInst){
        string code;
        /** 
         *code in here
        */
        return code;
    }
    pair<RegId, string> 
    CodeGen::callInst_gen(CallInst *callInst, RegId dstRegId){
        string code;
        /** 
         *code in here
        */
        return {dstRegId, code};
    }

    string CodeGen::instruction_gen(Instruction *instr, int *frame_size)
    {
        string code;
        string unkName = instr->getName();
        RegManager::RegId dstRegId = RegManager::RNONE;
        auto instrType = instr->getKind();
        pair<RegId, string> tmp;
#ifdef GEN_DEBUG
        char ch[50];
        int kKind = 0;
        for (int i = 0; i < 64; ++i)
            if ((1 << i) & instrType)
                kKind = i;
        sprintf(ch, "%s 's instrType: %llu", (instr->getName()).c_str(), kKind);
        std::cout << ch << "\n";
        std::cout.flush();
#endif // GEN_DEBUG
        switch (instrType)
        {
        //binary inst
        case Instruction::kAdd :
        case Instruction::kMul :
        case Instruction::kSub :
        {
            BinaryInst *bInst = dynamic_cast<BinaryInst *>(instr);
            //registers are used only for instruction operation, consider use which register (any one that is free for use)
            tmp = binaryInst_gen(bInst, RegManager::RANY);
            code += tmp.second;
            dstRegId = tmp.first;
            break;
        }
        case Instruction::kLoad :
        {
            LoadInst *ldInst = dynamic_cast<LoadInst *>(instr);
            tmp = loadInst_gen(ldInst, RegManager::RANY);
            // code += M_emitComment("load inst");
            code += tmp.second;
            dstRegId = tmp.first;
            break;
        }
        case Instruction::kStore :
        {
            StoreInst *stInst = dynamic_cast<StoreInst *>(instr);
            // code += M_emitComment("store inst");
            code += storeInst_gen(stInst);
            return code;
            break;
        }
        case Instruction::kAlloca:
        {
            AllocaInst *aInst = dynamic_cast<AllocaInst *>(instr);
            tmp = allocaInst_gen(aInst, RegManager::RANY, frame_size);
            // code += M_emitComment("alloca inst");
            code += tmp.second;
            dstRegId = tmp.first;
            break;
        }
        case Instruction::kReturn:
        {
            ReturnInst *retInst = dynamic_cast<ReturnInst *>(instr);
            code += M_emitComment("return inst");
            code += returnInst_gen(retInst);
            return code;
            break;
        }
        case Instruction::kCall:
        {
            CallInst *cInst = dynamic_cast<CallInst *>(instr);
            auto tmp = callInst_gen(cInst, RegManager::RANY);
            code += tmp.second;
            dstRegId = tmp.first;
            if(dstRegId == RegManager::R0) return code;
            break;
        }
        case Instruction::kBr:
        {
            UncondBrInst *ubInst = dynamic_cast<UncondBrInst *>(instr);
            code += uncondBrInst_gen(ubInst);
            return code;
            break;
        }
        case Instruction::kCondBr:
        {
            CondBrInst *cbInst = dynamic_cast<CondBrInst *>(instr);
            code += condBrInst_gen(cbInst);
            return code;
            break;
        }
        default:
        {
            code += "ERROR CODE : instruction " + \
                unkName + " is not implementation" + endl;
            break;
        }
        }
        if(!instr->getType()->isVoid()){
            code += storeRegToStack_gen(dstRegId, instr);
            // regm.freeReg(dstRegId);//TODO : code in here.
        }
        return code;
    }
    //
    string CodeGen::globaldata_gen(){
        string asmCode;
        /** 
         *code in here
        */
    //    for (auto a : compUnit->getGlobalValues()) {

    //    }
        return asmCode;
    }

    string CodeGen::literalPoolsCode_gen(Function *func){
        string code;
        /** 
         *code in here
        */
        return code;
    }

}//namespace backend
